package cn.hy.dubbo.limiter;

import cn.hy.dubbo.LimitConf;
import com.google.common.io.ByteStreams;
import org.apache.dubbo.common.URL;
import org.apache.dubbo.common.utils.StringUtils;
import org.apache.dubbo.rpc.Invocation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import javax.annotation.PostConstruct;
import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * @author zhaomu
 */
public class TokenBasedLimitStrategy implements LimitStrategy {
    private static final String logLabel = "[dubbo限流] ";
    private final Logger logger = LoggerFactory.getLogger(TokenBasedLimitStrategy.class);
    private static final String DUBBO_LIMIT_RUNTIME_PREFIX = "{dubbo_limit_runtime_}";
    private static final String DUBBO_LIMIT_TRIGGER_COUNT = "trigger_count";
    private DefaultRedisScript<Long> providerDefaultScript;
    private DefaultRedisScript<Long> consumerDefaultScript;
    private final StringRedisTemplate stringRedisTemplate;

    private final LimitConf limitConf;

    public TokenBasedLimitStrategy(StringRedisTemplate stringRedisTemplate,
                                   LimitConf limitConf) {
        this.stringRedisTemplate = stringRedisTemplate;
        this.limitConf = limitConf;
    }

    @PostConstruct
    public void init() {
        try {
            String providerLua = TokenBasedLimitStrategy.read("classpath:provider_rate_limiter.lua");
            assert providerLua != null;
            providerDefaultScript = new DefaultRedisScript<>(providerLua, Long.class);

            String consumerLua = TokenBasedLimitStrategy.read("classpath:consumer_rate_limiter.lua");
            assert consumerLua != null;
            consumerDefaultScript = new DefaultRedisScript<>(consumerLua, Long.class);
            logger.info(logLabel + "限流脚本缓存成功");
            // 订阅 变化， 清除缓存
            // stringRedisTemplate.execute()
        } catch (Exception e) {
            logger.error(logLabel + "初始化TokenBasedLimitStrategy报错", e);
        }


    }

    public JudgeResult isAllowable(URL url, Invocation invocation) {
        final JudgeResult rs = new JudgeResult();
        long startTime = System.currentTimeMillis();
        try {
            final String serviceKey = this.getServiceKey(url, invocation);
            if (!limitConf.isEnable()) {
                logger.debug(logLabel + "serviceKey=" + serviceKey + " enable=" + limitConf.isEnable() + "，跳过限流逻辑");
                return rs;
            }

            String side = url.getParameter("side");
            if (StringUtils.isEmpty(side)) {
                throw new Exception("没有side，跳过判断");
            }

            String limitRuntimeKey = getDubboLimitRunTimeKey(serviceKey);
            List<String> params = Arrays.asList(
                    String.valueOf((double) limitConf.getIntervalInMills() / (double) limitConf.getLimit()),
                    String.valueOf(System.currentTimeMillis()),
                    String.valueOf(limitConf.getBurstTokens()),
                    String.valueOf(limitConf.getLimit()),
                    String.valueOf(limitConf.getIntervalInMills()));
            List<String> keys = Collections.singletonList(limitRuntimeKey);
            Long result;
            if ("consumer".equals(side)) {
                result = stringRedisTemplate.execute(providerDefaultScript, keys, params);
            } else {
                result = stringRedisTemplate.execute(consumerDefaultScript, keys, params);
            }

            if (result == null || result <= 0L) {
                rs.setAllowable(false);
                rs.setErrorMsg("违背限流规则: " + limitConf);
                stringRedisTemplate.opsForHash().increment(limitRuntimeKey, TokenBasedLimitStrategy.DUBBO_LIMIT_TRIGGER_COUNT, 1L);
            }

        } catch (Exception e) {
            logger.error(logLabel + " 限流逻辑判断过程出错,本次默认通过，出错原因是", e);
        }

        long endTime = System.currentTimeMillis();
        logger.info(logLabel + " 判断完成，耗时:" + (endTime - startTime));
        return rs;
    }

    private String getDubboLimitRunTimeKey(String serviceKey) throws Exception {
        return DUBBO_LIMIT_RUNTIME_PREFIX + serviceKey;
    }

    /**
     * 生成key
     */
    private String getServiceKey(URL url, Invocation invocation) throws Exception {
        String application = url.getParameter("application");
        String interf = url.getParameter("interface");
        String version = url.getParameter("version");
        String methodName = invocation.getMethodName();
        Class<?>[] classes = invocation.getParameterTypes();
        StringBuffer sb = new StringBuffer();
        sb.append(application).append("_");
        sb.append(interf).append("_");
        sb.append(version).append("_");
        sb.append(methodName).append("_");
        if (classes != null && classes.length > 0) {
            Class<?>[] paramTypesArr = classes;
            int length = classes.length;

            for (int i = 0; i < length; ++i) {
                Class cls = paramTypesArr[i];
                sb.append(cls.getName()).append("_");
            }
        }
        String rs = sb.toString();
        if (rs.endsWith("_")) {
            rs = rs.substring(0, rs.length() - 1);
        }
        String side = url.getParameter("side");
        if (StringUtils.isEmpty(side)) {
            throw new Exception("不确定side，无法决定是provider还是consumer");
        } else if ("consumer".equals(side)) {
            String host = url.getHost();
            String pid = url.getParameter("pid");
            String consumer = host + ":" + pid;
            return rs + "_" + side + "_" + consumer;
        } else {
            return rs + "_" + side;
        }
    }

    /**
     * 读取配置文件
     */
    private static String read(String luaScriptPath) throws IOException {
        PathMatchingResourcePatternResolver patternResolver = new PathMatchingResourcePatternResolver();
        Resource[] resources = patternResolver.getResources(luaScriptPath);
        if (resources.length > 0) {
            InputStream inputStream = resources[0].getInputStream();
            byte[] scriptBytes = ByteStreams.toByteArray(inputStream);
            inputStream.close();
            return new String(scriptBytes);
        } else {
            return null;
        }
    }
}
